# -*- coding: utf-8 -*-
import numpy as np 
import pandas as pd
from tqdm import tqdm  
from torch import nn 
import matplotlib 
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import os
from typing import Generator
import torch
import fnmatch
import sys
sys.path.append("/gpfs/scratch/chgwang/XI/Scripts/Refactoring_1/MLModel")
import lstm_net # type: ignore
import lstm_train # type: ignore

def pltFile(path:str,
            model:torch.nn.Module,
            oppsite=False,):
    with open(path, mode="r") as f:
        while True:
            line = f.readline()
            if "SampleRate" in line:
                break
    line = line.split(",")
    # read the frequency
    source_freq = float(line[1].strip())
    # our sample rate
    resample_freq = 1e4
    freq_times = int(source_freq / resample_freq)
    # get the labels
    path_splited = path.split("/")
    # start with . means is the hidden data
    if path_splited[-1][0] == ".":
        # print(path)
        return
    labels = []
    if path_splited[-1][0] == "0":
        labels.append("0")
    elif not path_splited[-1][0].isdigit():
        labels.append("0")
    else:
        labels.append(path_splited[-1][0])
        # isdigit is a function for str.m
        if path_splited[-1][1].isdigit():
            labels.append(path_splited[-1][1])
    # convert data to int
    labels = np.array(labels, dtype=np.int16)
    sour_data = np.loadtxt(path, skiprows=16, delimiter=",", usecols=range(1,4))
    assert sour_data.shape[0] == 125000
    # take the oppsite number.
    if oppsite:
        sour_data = -sour_data
    modeled_data = sour_data[::freq_times, :3]
    # modeled_data = np.transpose(modeled_data)
    delimiter = int(modeled_data.shape[0] / 2)
    target_label_arr = np.zeros((6, modeled_data.shape[0]))
    target_label_arr = np.transpose(target_label_arr)
    for label in labels:
        if label != 0:
            label = label - 1
            # using broadcast property
            target_label_arr[delimiter:,label] = 1
    model.eval()
    modeled_data = np.expand_dims(modeled_data, axis=0)
    inp_seq = torch.from_numpy(modeled_data)
    inp_seq = inp_seq.float()
    target_label_arr = np.expand_dims(target_label_arr, axis=0)
    target_label_arr = torch.from_numpy(target_label_arr)
    target_label_arr = target_label_arr.float()
    modeled_label_arr, _ = model(inp_seq)

    # plot the modeling result.
    ## prepare the data.
    modeled_label_arr = modeled_label_arr.detach().numpy()
    modeled_label_arr = np.squeeze(modeled_label_arr)
    target_label_arr = target_label_arr.numpy()
    target_label_arr = np.squeeze(target_label_arr)
    modeled_data = np.squeeze(modeled_data)
    ## plot action.
    markers = ["b", "r:", "y--"]
    fig_path = path.replace(".csv", ".png")
    fig, axes = plt.subplots(nrows=len(labels)+1, ncols=1,
                sharex=True)
    axes[0].plot(modeled_data[:,0], markers[0], label="ia")
    axes[0].plot(modeled_data[:,1], markers[1], label="ib")
    axes[0].plot(modeled_data[:,2], markers[2], label="ic")
    axes[0].legend(loc=1)
    x_labeled = np.arange(199, modeled_data.shape[0])
    for index, label in enumerate(labels):
        if label != 0:
            label = label - 1
            axes[index+1].plot(x_labeled,
                            target_label_arr[199:,label], 
                            markers[0], label="target")
            axes[index+1].plot(x_labeled,
                            modeled_label_arr[199:,label], 
                            markers[1], label="modeling")
            axes[index+1].legend(loc=1)
        else:
            axes[index+1].plot(x_labeled,
                            target_label_arr[199:,label], 
                            markers[0], label="target")
            axes[index+1].plot(x_labeled,
                            target_label_arr[199:,label], 
                            markers[0], label="modeling")
            axes[index+1].legend(loc=1)
    fig.savefig(fig_path)
    plt.close()

def iter_image(path_list:list, model:nn.Module):
    for path_num, path in enumerate(path_list):
        for file in lstm_train.retrieve_files(path):
            if fnmatch.fnmatch(file, "*.csv"):
                if path_num == 0:
                    pltFile(file, model, oppsite=True)
                else:
                    pltFile(file, model)

if __name__ == "__main__":
    # oppsite
    path_0 = "/gpfs/scratch/chgwang/XI/data/论文展示的数据/1.整流部分---实验二"
    # non-oppsite 
    path_1 = "/gpfs/scratch/chgwang/XI/data/论文展示的数据/2.逆变部分---第三次实验"
    path_2 = "/gpfs/scratch/chgwang/XI/data/论文展示的数据/3.特殊情况" 
    path_list = [path_0, path_1, path_2]
    # model setting
    input_dim = 3
    hidden_size = 1024
    num_layers = 4
    out_dim = 6
    model = lstm_net.lstm_net(input_size=input_dim,
    hidden_size=hidden_size, num_layers=num_layers, out_dim=out_dim)
    model_path = "/gpfs/scratch/chgwang/XI/DataBase/Model_lstm_mse/PN-19-0.193.pt"
    trained_dict = torch.load(model_path, map_location="cpu")
    if "model_state_dict" in trained_dict:
        trained_dict = trained_dict["model_state_dict"]
    model.load_state_dict(trained_dict, strict=True)
    # test_path = "/gpfs/scratch/chgwang/XI/data/论文展示的数据/1.整流部分---实验二/12.csv"
    iter_image(path_list, model)